import torch
import torchvision.transforms as transforms
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np

# 加载模型定义（与训练时一致）
class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = torch.nn.Conv2d(1, 32, 3, 1)
        self.conv2 = torch.nn.Conv2d(32, 64, 3, 1)
        self.fc1 = torch.nn.Linear(12*12*64, 128)
        self.fc2 = torch.nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = torch.nn.ReLU()(x)
        x = self.conv2(x)
        x = torch.nn.ReLU()(x)
        x = torch.nn.MaxPool2d(2)(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = torch.nn.ReLU()(x)
        x = self.fc2(x)
        return x

# 初始化模型并加载权重
net = Net()
net.load_state_dict(torch.load('mnist_model.pth'))
net.eval()  # 将模型设置为评估模式

# 图像预处理
transform = transforms.Compose([
    transforms.Grayscale(num_output_channels=1),  # 将图像转换为灰度图像
    transforms.Resize((28, 28)),                  # 调整大小到28x28
    transforms.ToTensor(),                        # 转换为Tensor
    transforms.Normalize((0.5,), (0.5,))          # 归一化
])

# 加载本地图像文件
image_path = '3.png'  # 替换为你的图像路径
image = Image.open(image_path).convert('L')  # 确保图像是灰度图
image = transform(image)
image = image.unsqueeze(0)  # 添加批次维度

# 显示原始图像
def imshow(img):
    img = img / 2 + 0.5  # 去归一化
    npimg = img.numpy()
    if npimg.shape[0] == 1:
        npimg = npimg[0]  # 去除批次维度
    plt.imshow(npimg, cmap='gray')
    plt.show()

# 显示输入图像
imshow(image.squeeze())

# 进行预测
with torch.no_grad():
    output = net(image)
_, predicted = torch.max(output, 1)

print('Predicted: ', predicted.item())